import asyncio
import json
import random
import uuid
from datetime import datetime
from typing import Dict, Set, Any
import logging
from aiohttp import web, WSMsgType
from user import register, login, users, generate_user_info
from friends import add_friend, remove_friend
from private_chat import handle_private_chat_aiohttp
import aiohttp_cors

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 存储连接的客户端
clients: Dict[web.WebSocketResponse, Dict[str, Any]] = {}
client_id_counter = 0
user_id_to_ws = {}

# 获取在线用户列表
def get_online_users() -> list:
    users_list = []
    for ws, user_data in clients.items():
        users_list.append({
            'id': user_data['id'],
            'username': user_data['username'],
            'avatar_color': user_data.get('avatar_color', '#666'),
            'avatar_text': user_data.get('avatar_text', user_data['username'][:2])
        })
    return users_list

# 广播消息给所有客户端
async def broadcast(message: Dict[str, Any], exclude_client: web.WebSocketResponse = None):
    disconnected_clients = set()
    for client in list(clients.keys()):
        if client != exclude_client:
            try:
                await client.send_str(json.dumps(message, ensure_ascii=False, default=str))
            except Exception as e:
                logger.error(f"发送消息时出错: {e}")
                disconnected_clients.add(client)
    for client in disconnected_clients:
        if client in clients:
            del clients[client]

# 处理客户端断开连接
async def handle_client_disconnect(ws: web.WebSocketResponse):
    user_data = clients.get(ws)
    if user_data:
        logger.info(f"用户 {user_data['username']} (ID: {user_data['id']}) 已断开连接")
        del clients[ws]
        await broadcast({
            'type': 'userLeft',
            'data': {
                'userId': user_data['id'],
                'username': user_data['username'],
                'message': f"{user_data['username']} 离开了聊天室",
                'timestamp': datetime.now()
            }
        })
        await broadcast({
            'type': 'userList',
            'data': get_online_users()
        })

# WebSocket 路由
async def websocket_handler(request):
    global client_id_counter
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    client_id_counter += 1
    client_id = client_id_counter
    user_info = generate_user_info()
    clients[ws] = {
        'id': client_id,
        'username': user_info['username'],
        'avatar_color': user_info['avatar_color'],
        'avatar_text': user_info['avatar_text'],
        'join_time': datetime.now()
    }
    logger.info(f"用户 {user_info['username']} (ID: {client_id}) 已连接")
    # 发送欢迎消息
    await ws.send_str(json.dumps({
        'type': 'welcome',
        'data': {
            'userId': client_id,
            'username': user_info['username'],
            'avatar_color': user_info['avatar_color'],
            'avatar_text': user_info['avatar_text'],
            'message': f"欢迎 {user_info['username']} 加入聊天室！"
        }
    }, ensure_ascii=False, default=str))
    # 广播用户加入
    await broadcast({
        'type': 'userJoined',
        'data': {
            'userId': client_id,
            'username': user_info['username'],
            'avatar_color': user_info['avatar_color'],
            'avatar_text': user_info['avatar_text'],
            'message': f"{user_info['username']} 加入了聊天室",
            'timestamp': datetime.now()
        }
    }, ws)
    # 广播在线用户列表
    await broadcast({
        'type': 'userList',
        'data': get_online_users()
    })
    try:
        async for msg in ws:
            if msg.type == WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    msg_type = data.get('type')
                    if msg_type == 'bind':
                        user_id = data['data']['user_id']
                        clients[ws]['user_id'] = user_id
                        user_id_to_ws[user_id] = ws
                        for u in users.values():
                            if u['id'] == user_id:
                                clients[ws]['username'] = u['username']
                                clients[ws]['avatar_color'] = u['avatar_color']
                                clients[ws]['avatar_text'] = u['avatar_text']
                        logger.info(f"WebSocket绑定user_id: {user_id}")
                    elif msg_type == 'chat':
                        user_data = clients[ws]
                        message_id = str(uuid.uuid4())
                        chat_message = {
                            'type': 'chat',
                            'data': {
                                'messageId': message_id,
                                'userId': user_data['id'],
                                'username': user_data['username'],
                                'avatar_color': user_data.get('avatar_color', '#666'),
                                'avatar_text': user_data.get('avatar_text', user_data['username'][:2]),
                                'text': data['data']['text'],
                                'timestamp': datetime.now()
                            }
                        }
                        logger.info(f"{user_data['username']}: {data['data']['text']}")
                        await broadcast(chat_message)
                    elif msg_type == 'typing':
                        typing_data = clients[ws]
                        await broadcast({
                            'type': 'typing',
                            'data': {
                                'userId': typing_data['id'],
                                'username': typing_data['username'],
                                'isTyping': data['data']['isTyping']
                            }
                        }, ws)
                    elif msg_type == 'image':
                        image_data = clients[ws]
                        message_id = str(uuid.uuid4())
                        image_message = {
                            'type': 'image',
                            'data': {
                                'messageId': message_id,
                                'userId': image_data['id'],
                                'username': image_data['username'],
                                'avatar_color': image_data.get('avatar_color', '#666'),
                                'avatar_text': image_data.get('avatar_text', image_data['username'][:2]),
                                'imageId': data['data']['imageId'],
                                'fileName': data['data']['fileName'],
                                'imageData': data['data']['imageData'],
                                'size': data['data']['size'],
                                'timestamp': datetime.now()
                            }
                        }
                        await broadcast(image_message)
                    elif msg_type == 'private_chat':
                        await handle_private_chat_aiohttp(data, ws, clients, user_id_to_ws)
                    else:
                        logger.warning(f"未知消息类型: {msg_type}")
                except Exception as e:
                    logger.error(f"处理消息时出错: {e}")
            elif msg.type == WSMsgType.ERROR:
                logger.error(f'ws connection closed with exception {ws.exception()}')
    finally:
        await handle_client_disconnect(ws)
        # 清理user_id_to_ws
        if 'user_id' in clients[ws]:
            user_id_to_ws.pop(clients[ws]['user_id'], None)
    return ws

# 新增欢迎页
async def index(request):
    return web.Response(text='IM 聊天后端服务已启动。')

# 启动 aiohttp 服务
async def main():
    app = web.Application()
    # 欢迎页
    app.router.add_get('/', index)
    # HTTP接口
    app.router.add_post('/register', register)
    app.router.add_post('/login', login)
    app.router.add_post('/add_friend', add_friend)
    app.router.add_post('/remove_friend', remove_friend)
    # WebSocket接口
    app.router.add_get('/ws', websocket_handler)

    # CORS支持
    cors = aiohttp_cors.setup(app, defaults={
        "*": aiohttp_cors.ResourceOptions(
            allow_credentials=True,
            expose_headers="*",
            allow_headers="*",
        )
    })
    for route in list(app.router.routes()):
        cors.add(route)

    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, 'localhost', 3000)
    await site.start()
    logger.info(f"HTTP+WebSocket服务运行在端口 3000")
    logger.info(f"WebSocket地址: ws://localhost:3000/ws")
    logger.info("按 Ctrl+C 停止服务器")
    while True:
        await asyncio.sleep(3600)

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("正在关闭服务器...")
        logger.info("WebSocket服务器已关闭") 